import base64
import time
from enum import StrEnum
from typing import Literal, TypedDict

from PIL import Image

from anthropic.types.beta import BetaToolComputerUse20241022Param

from .base import BaseAnthropicTool, ToolError, ToolResult
from .screen_capture_android import get_android_screenshot
import requests
import re

import subprocess

OUTPUT_DIR = "./tmp/outputs"

TYPING_DELAY_MS = 12
TYPING_GROUP_SIZE = 50

Action = Literal[
    "type",
    "tap",
    "longpress",
    "swipe_up",
    "swipe_down",
    "swipe_left",
    "swipe_right",
    "screenshot",
    "hover",
    "wait",
    "back",
]


class Resolution(TypedDict):
    width: int
    height: int

# TODO 仅仅测了一台手机的720x1440的情况
MAX_SCALING_TARGETS: dict[str, Resolution] = {
    # "XGA": Resolution(width=1024, height=768),  # 4:3
    # "WXGA": Resolution(width=1280, height=800),  # 16:10
    # "FWXGA": Resolution(width=1366, height=768),  # ~16:9
    "FHD": Resolution(width=720, height=1280), # 9:16
    "FHD2": Resolution(width=720, height=1440), # 9:18
    "FHD3": Resolution(width=720, height=1600), # 9:20
}


class ScalingSource(StrEnum):
    COMPUTER = "computer"
    API = "api"


class AndroidToolOptions(TypedDict):
    display_height_px: int
    display_width_px: int
    # display_number: int | None


def chunks(s: str, chunk_size: int) -> list[str]:
    return [s[i : i + chunk_size] for i in range(0, len(s), chunk_size)]

class AndroidTool(BaseAnthropicTool):
    """
    A tool that allows the agent to interact with the screen, keyboard, and mouse of the current computer.
    Adapted for Windows using 'pyautogui'.
    """

    name: Literal["computer"] = "computer"
    api_type: Literal["computer_20241022"] = "computer_20241022"
    width: int
    height: int
    display_num: int | None

    _screenshot_delay = 2.0
    _scaling_enabled = True

    @property
    def options(self) -> AndroidToolOptions:
        width, height = self.scale_coordinates(
            ScalingSource.COMPUTER, self.width, self.height
        )
        return {
            "display_width_px": width,
            "display_height_px": height,
            # "display_number": self.display_num,
        }

    def to_params(self) -> BetaToolComputerUse20241022Param:
        return {"name": self.name, "type": self.api_type, **self.options}

    def __init__(self, is_scaling: bool = False):
        super().__init__()

        # Get screen width and height using Windows command
        self.display_num = None
        self.offset_x = 0
        self.offset_y = 0
        self.is_scaling = is_scaling
        self.width, self.height = self.get_screen_size()
        print(f"screen size: {self.width}, {self.height}")



    async def __call__(
        self,
        *,
        action: Action,
        text: str | None = None,
        coordinate: tuple[int, int] | None = None,
        **kwargs,
    ):
        print(f"action: {action}, text: {text}, coordinate: {coordinate}, is_scaling: {self.is_scaling}")

        if action == "type":
            if text is None:
                raise ToolError(f"text is required for {action}")
            if coordinate is not None:
                raise ToolError(f"coordinate is not accepted for {action}")
            if not isinstance(text, str):
                raise ToolError(output=f"{text} must be a string")

            # default click before type: check if this is needed
            # self.send_to_vm("pyautogui.click()")
            # self.send_to_vm(f"pyautogui.typewrite('{text}', interval={TYPING_DELAY_MS / 1000})")
            # self.send_to_vm("pyautogui.press('enter')")
            self._execute_adb(['shell', 'input', 'text', f'"{text}"'])
            screenshot_base64 = (await self.screenshot()).base64_image
            return ToolResult(output=text, base64_image=screenshot_base64)

        if action in (
            "tap",
            "longpress"
        ):
            if text is not None:
                raise ToolError(f"text is not accepted for {action}")
            if coordinate is None:
                raise ToolError(f"coordinate is required for {action}")

            
            elif action == 'tap':
                x, y = self.scale_coordinates(ScalingSource.COMPUTER, coordinate[0], coordinate[1])
                return self._execute_adb(['shell', 'input', 'tap', str(x), str(y)])
            elif action == 'longpress':
                x, y = self.scale_coordinates(ScalingSource.COMPUTER, coordinate[0], coordinate[1])
                return self._execute_adb([
                    'shell', 'input', 'swipe',
                    str(x), str(y), str(x), str(y), '1500'
                ])
                
        if action in ("swipe_up", "swipe_down", "swipe_left", "swipe_right"):
            start_x, start_y, end_x, end_y = self._get_swipe_coords(action)
            return self._execute_adb([
                'shell', 'input', 'swipe',
                str(start_x), str(start_y), str(end_x), str(end_y)
            ])
            # return ToolResult(output=f"Performed {action}")
        if action == "back":
            return self._execute_adb(['shell', 'input', 'keyevent', '4'])
        if action == "screenshot":
            return await self.screenshot()
        if action == "wait":
            time.sleep(1)
            return ToolResult(output=f"Performed {action}")
        raise ToolError(f"Invalid action: {action}")
    def _get_swipe_coords(self, action: str) -> tuple[int, int, int, int]:
        """生成上半屏滑动坐标（避开软键盘）"""
        center_x = self.width // 2
        upper_quarter = self.height // 4
        
        return {
            'swipe_up': (center_x, upper_quarter*3, center_x, upper_quarter),
            'swipe_down': (center_x, upper_quarter, center_x, upper_quarter*3),
            'swipe_left': (self.width*3//4, upper_quarter*2, self.width//4, upper_quarter*2),
            'swipe_right': (self.width//4, upper_quarter*2, self.width*3//4, upper_quarter*2)
        }[action]
    
    def _execute_adb(self, command: list) -> ToolResult:
        """执行ADB命令的核心方法"""
        print('adb', command)
        try:
            result = subprocess.run(
                ['adb'] + command,
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                timeout=10
            )
            return ToolResult(output=result.stdout.decode())
        except subprocess.CalledProcessError as e:
            # self._action_history.append({
            #     'command': command,
            #     'success': False,
            #     'error': e.stderr.decode()
            # })
            return ToolError(f'ADB error: {e.stderr.decode()}')

    async def screenshot(self):
        if not hasattr(self, 'target_dimension'):
            screenshot = self.padding_image(screenshot)
            self.target_dimension = MAX_SCALING_TARGETS["FHD2"]
        width, height = self.target_dimension["width"], self.target_dimension["height"]
        screenshot, path = get_android_screenshot(resize=True, target_width=width, target_height=height)
        time.sleep(0.7) # avoid async error as actions take time to complete
        return ToolResult(base64_image=base64.b64encode(path.read_bytes()).decode())

    def padding_image(self, screenshot):
        """Pad the screenshot to 9:18 aspect ratio, when the aspect ratio is not 9:18."""
        _, height = screenshot.size
        new_width = height * 9 // 18

        padding_image = Image.new("RGB", (new_width, height), (255, 255, 255))
        # padding to top left
        padding_image.paste(screenshot, (0, 0))
        return padding_image

    def scale_coordinates(self, source: ScalingSource, x: int, y: int):
        """Scale coordinates to a target maximum resolution."""
        if not self._scaling_enabled:
            return x, y
        ratio = self.width / self.height
        target_dimension = None

        for target_name, dimension in MAX_SCALING_TARGETS.items():
            # allow some error in the aspect ratio - not ratios are exactly 16:9
            if abs(dimension["width"] / dimension["height"] - ratio) < 0.02:
                if dimension["width"] < self.width:
                    target_dimension = dimension
                    self.target_dimension = target_dimension
                    # print(f"target_dimension: {target_dimension}")
                break

        if target_dimension is None:
            # TODO: currently we force the target to be WXGA (16:10), when it cannot find a match
            target_dimension = MAX_SCALING_TARGETS["FHD2"]
            self.target_dimension = MAX_SCALING_TARGETS["FHD2"]

        # should be less than 1
        x_scaling_factor = target_dimension["width"] / self.width
        y_scaling_factor = target_dimension["height"] / self.height
        if source == ScalingSource.API:
            if x > self.width or y > self.height:
                raise ToolError(f"Coordinates {x}, {y} are out of bounds")
            # scale up
            return round(x / x_scaling_factor), round(y / y_scaling_factor)
        # scale down
        return round(x * x_scaling_factor), round(y * y_scaling_factor)
    

    def get_screen_size(self):
        """通过ADB获取实时屏幕分辨率(优先返回Override值)"""
        try:
            output = subprocess.check_output(
                ['adb', 'shell', 'wm', 'size'],
                stderr=subprocess.STDOUT,
                timeout=5
            ).decode()
            
            # 优先匹配Override尺寸
            override_match = re.search(r'Override size:\s*(\d+)x(\d+)', output)
            if override_match:
                return tuple(map(int, override_match.groups()))
            
            # 回退到物理分辨率
            physical_match = re.search(r'Physical size:\s*(\d+)x(\d+)', output)
            if physical_match:
                return tuple(map(int, physical_match.groups()))
            
            return (1080, 2160)  # 18:9默认分辨率[3,5](@ref)
        
        except subprocess.TimeoutExpired as e:
            raise ToolError(f"ADB连接超时: {str(e)}")